package com.kusej.dbdocexport.utils;

import com.intellij.database.access.ConnectionProvider;
import com.intellij.database.dataSource.DatabaseConnection;
import com.intellij.database.dataSource.connection.DGDepartment;
import com.intellij.database.psi.DbDataSource;
import com.intellij.database.remote.jdbc.RemotePreparedStatement;
import com.intellij.database.remote.jdbc.RemoteResultSet;
import com.intellij.database.util.GuardedRef;
import com.kusej.dbdocexport.model.TableColumnInfo;
import com.kusej.dbdocexport.model.TableIndexInfo;
import com.kusej.dbdocexport.model.TableInfo;
import org.apache.commons.compress.utils.Lists;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.stream.Collectors;

public class DBExecutorUtils {
    /**
     * 查询全部schema
     */
    private static final String selectSchemaSql = "select schema_name from information_schema.schemata ";
    /**
     * 查询schema下全部表
     */
    private static final String selectTableSql = "select table_name,table_comment from information_schema.tables where table_schema= ? order by table_name";
    /**
     * 查询 Schema 下全部索引
     */
    private static final String selectIndexSql = "select table_name,index_name,group_concat(column_name order by seq_in_index) as column_name,max(non_unique) as non_unique,max(index_comment) as index_comment from information_schema.statistics where table_schema = ? group by table_name,index_name order by table_name,if(index_name='PRIMARY',0,1),index_name";
    /**
     * 查询 Schema 下全部元数据
     */
    private static final String selectColumnSql = "select table_name,column_name,column_type,column_default,is_nullable,column_comment from information_schema.columns where table_schema = ? order by table_name,ordinal_position";



    public static List<String> getSchemas(DbDataSource dbDataSource){
//        if (CollectionUtils.isNotEmpty(schemaNames)){
//            String join = Joiner.on("','").join(schemaNames).concat("')");
//            selectSql = "select schema_name from information_schema.schemata where schema_name in ('"+join;
//        }
        List<String> schemaList = Lists.newArrayList();
        final ConnectionProvider provider = ConnectionProvider.forElement(dbDataSource, DGDepartment.INTROSPECTION);
        try (GuardedRef<DatabaseConnection> databaseConnectionGuardedRef = provider.acquireImpl()){
            RemotePreparedStatement remotePreparedStatement = databaseConnectionGuardedRef.get().getRemoteConnection().prepareStatement(selectSchemaSql);
            RemoteResultSet remoteResultSet = remotePreparedStatement.executeQuery();
            while (remoteResultSet.next()) {
                String schemaName = remoteResultSet.getString("schema_name");
                schemaList.add(schemaName);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return schemaList;
    }


    public static List<TableInfo> combineAllInfo(String dbSchema, DbDataSource dbDataSource){
        List<TableInfo> tableInfoList = getTableInfo(dbSchema,dbDataSource);
        List<TableIndexInfo> tableIndexInfoList = getTableIndexInfo(dbSchema,dbDataSource);
        List<TableColumnInfo> tableColumnInfoList = getTableColumnInfo(dbSchema,dbDataSource);

        LinkedHashMap<String, List<TableIndexInfo>> indexInfoMap = tableIndexInfoList.stream()
                .collect(Collectors.groupingBy(TableIndexInfo::getTableName,LinkedHashMap::new, Collectors.toList()));
        LinkedHashMap<String, List<TableColumnInfo>> columnInfoMap = tableColumnInfoList.stream()
                .collect(Collectors.groupingBy(TableColumnInfo::getTableName,LinkedHashMap::new, Collectors.toList()));

        for (TableInfo tableInfo : tableInfoList) {
            String tableName = tableInfo.getTableName();
            List<TableIndexInfo> oneTableIndexInfo = indexInfoMap.get(tableName);
            if (oneTableIndexInfo != null){
                int index = 1;
                for (TableIndexInfo tableIndexInfo : oneTableIndexInfo) {
                    tableIndexInfo.setIndex(index);
                    index++;
                }
                tableInfo.setIndexInfoList(oneTableIndexInfo);
            }else {
                tableInfo.setIndexInfoList(Collections.emptyList());
            }


            List<TableColumnInfo> oneTableColumnInfo = columnInfoMap.get(tableName);
            if (oneTableColumnInfo != null){
                int index = 1;
                for (TableColumnInfo tableColumnInfo : oneTableColumnInfo) {
                    tableColumnInfo.setIndex(index);
                    index++;
                }
                tableInfo.setColumnInfoList(oneTableColumnInfo);
            }else {
                tableInfo.setColumnInfoList(Collections.emptyList());
            }


        }

        return tableInfoList;
    }

    private static List<TableColumnInfo> getTableColumnInfo(String dbSchema, DbDataSource dbDataSource) {
        List<TableColumnInfo> tableColumnInfoList = new ArrayList<>();
        final ConnectionProvider provider = ConnectionProvider.forElement(dbDataSource, DGDepartment.INTROSPECTION);
        try (GuardedRef<DatabaseConnection> databaseConnectionGuardedRef = provider.acquireImpl()){
            RemotePreparedStatement remotePreparedStatement = databaseConnectionGuardedRef.get().getRemoteConnection().prepareStatement(selectColumnSql);
            remotePreparedStatement.setString(1,dbSchema);
            RemoteResultSet rs = remotePreparedStatement.executeQuery();
            while (rs.next()) {
                String tableName = rs.getString("table_name");
                String columnName = rs.getString("column_name");
                String columnType = rs.getString("column_type");
                String columnDefault = rs.getString("column_default");
                String columnComment = rs.getString("column_comment");
                String isNullable = rs.getString("is_nullable");
                TableColumnInfo tableColumnInfo = new TableColumnInfo();
                tableColumnInfo.setTableName(tableName);
                tableColumnInfo.setColumnName(columnName);
                tableColumnInfo.setColumnType(columnType);
                tableColumnInfo.setColumnDefaultValue(String.valueOf(columnDefault));
                tableColumnInfo.setColumnComment(columnComment);
                tableColumnInfo.setIsNullable("YES".equals(isNullable));
                tableColumnInfoList.add(tableColumnInfo);
            }
            return tableColumnInfoList;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    private static List<TableIndexInfo> getTableIndexInfo(String dbSchema, DbDataSource dbDataSource) {
        List<TableIndexInfo> tableIndexInfoList = new ArrayList<>();
        final ConnectionProvider provider = ConnectionProvider.forElement(dbDataSource, DGDepartment.INTROSPECTION);
        try (GuardedRef<DatabaseConnection> databaseConnectionGuardedRef = provider.acquireImpl()){
            RemotePreparedStatement remotePreparedStatement = databaseConnectionGuardedRef.get().getRemoteConnection().prepareStatement(selectIndexSql);
            remotePreparedStatement.setString(1,dbSchema);
            RemoteResultSet rs = remotePreparedStatement.executeQuery();
            while (rs.next()) {
                String tableName = rs.getString("table_name");
                String indexName = rs.getString("index_name");
                String columnName = rs.getString("column_name");
                boolean nonUnique = rs.getBoolean("non_unique");
                String indexComment = rs.getString("index_comment");
                TableIndexInfo tableIndexInfo = new TableIndexInfo();
                tableIndexInfo.setTableName(tableName);
                tableIndexInfo.setIndexName(indexName);
                tableIndexInfo.setRelatedColumn(columnName);
                tableIndexInfo.setIsUnique(!nonUnique);
                tableIndexInfo.setIndexComment(indexComment);
                tableIndexInfoList.add(tableIndexInfo);
            }
            return tableIndexInfoList;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    private static List<TableInfo> getTableInfo(String dbSchema, DbDataSource dbDataSource) {
        List<TableInfo> tableInfoList = new ArrayList<>();
        final ConnectionProvider provider = ConnectionProvider.forElement(dbDataSource, DGDepartment.INTROSPECTION);
        try (GuardedRef<DatabaseConnection> databaseConnectionGuardedRef = provider.acquireImpl()){
            RemotePreparedStatement remotePreparedStatement = databaseConnectionGuardedRef.get().getRemoteConnection().prepareStatement(selectTableSql);
            remotePreparedStatement.setString(1,dbSchema);
            RemoteResultSet rs = remotePreparedStatement.executeQuery();
            int index = 1;
            while (rs.next()) {
                String tableName = rs.getString("table_name");
                String tableComment = rs.getString("table_comment");
                TableInfo tableInfo = new TableInfo();
                tableInfo.setTableName(tableName);
                tableInfo.setTableComment(tableComment);
                tableInfo.setIndex(index);
                tableInfoList.add(tableInfo);
                index++;
            }
            return tableInfoList;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }
}
